/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2019-2022. All rights reserved.
 * Description: 集合通信算子信息库头文件
 * Author: lilianlin
 * Create: 2019-11-28
 */

#ifndef HVD_OPS_KERNEL_INFO_STORE_H
#define HVD_OPS_KERNEL_INFO_STORE_H

#include "ops_kernel_info_store_base.h"
#include "op_hvd_comm.h"
#include "runtime/event.h"
#include "runtime/stream.h"

namespace hccl {
extern bool g_hvdMode;
// Ge适配的类
class HvdOpsKernelInfoStore : public HCCLOpsKernelInfoStore {
public:
    HvdOpsKernelInfoStore();
    ~HvdOpsKernelInfoStore() override;
    // load task for op
    ge::Status LoadTask(ge::GETaskInfo &task) override;
    ge::Status UnloadTask(ge::GETaskInfo &task) override;
private:
    HcclResult SetCustomKernelInfo(ge::OpInfo &opinfo, std::map<string, ge::OpInfo> &infos) const override;
    HcclResult GetSupportedOP(std::vector<std::string> &hcclSupportOp) const override;
    HcclResult GetDataTypeFromTaskInfo(const ge::GETaskInfo &task, HcclDataType &dataType) const override;
    HcclResult HCCLOpsKernel(const ge::GETaskInfo &task, const std::string &sCollectiveType);
    HcclResult HvdBroadcastOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType);
    HcclResult HvdAllReduceOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType);
    HcclResult HvdAllGatherOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType);
    HcclResult HvdWaitOpKernel(const ge::GETaskInfo &task, const string &sCollectiveType);
    HcclResult GetEventIDFromTaskInfo(const ge::GETaskInfo &task, u32 &eventID);
    HcclResult GetEvent(const u32 &eventID, rtEvent_t &event);
    HcclResult DestroyEvent();
    std::map<u32, rtEvent_t> eventMap_;
    std::vector<HvdCallBackInfo *> callbackVec_;
};
}
#endif // GE_OPS_KERNEL_INFO_H